#!/usr/bin/env python

import sys
import onnx
import argparse
from pathlib import Path

SEP = '*'
SEP_LEN = 80

PROG = Path(sys.argv[0]).name

def __sep():
    print (SEP * SEP_LEN)

def __parse_args():
    parser = argparse.ArgumentParser(prog=PROG, description='输出指定onnx算子的说明文档')
    parser.add_argument('--operator', '-op', type=str, help='算子名称，不区分大小写。缺省列出所有算子名称。')
    parser.add_argument('--attribute', '-att', action='store_true', help='列出attribute的描述。')
    return parser.parse_args()

def __list_ops():
    schemas = onnx.defs.get_all_schemas()

    # 打印操作列表
    for idx, s in enumerate(schemas):
        print(f'{idx}: {s.name}')

def __get_op(opName):
    for s in onnx.defs.get_all_schemas():
        if s.name.lower() == opName.lower():
            return s
    return None

def __print_op_doc(args):
    op = __get_op(args.operator)
    if op is None:
        raise Exception (f'Not found operator {args.operator}')
    print(f'{op.name} doc')
    __sep()
    print(op.doc)
    if args.attribute:
        __print_attributes(op)

def __print_attributes(op):
    print ('\nAttributes') 
    __sep()
    for attName, att in op.attributes.items():
        print(f'\n{attName}')
        print(f'type        :{att.type}')
        print(f'required    :{att.required}')
        # print(f'default     :{att.default_value}')
        print(f'description :{att.description}')

def main():
    args = __parse_args()
    if args.operator is None:
        __list_ops()
        return
    __print_op_doc(args)

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        print(e)
        sys.exit(1)
